package org.etnaframework.module.base.utils;

import org.etnaframework.module.base.config.ConstantKeys;
import org.slf4j.MDC;

import java.util.concurrent.Callable;
import java.util.concurrent.Executor;

/**
 * 线程链路追踪辅助工具
 *
 * @author jasyaf
 * @since 2025-03-06
 */
public class TraceLogUtils {

    /**
     * 将当前线程的traceId传递给指定的executor执行，这样在后台线程执行的任务也能追踪到请求链路
     */
    public static void executeWithTraceId(Executor executor, Runnable runnable) {
        String currentTraceId = MDC.get(ConstantKeys.TAG_TRACE_ID);
        executor.execute(() -> {
            try {
                MDC.put(ConstantKeys.TAG_TRACE_ID, currentTraceId);
                runnable.run();
            } finally {
                MDC.clear();
            }
        });
    }

    /**
     * 使用指定的traceId标记任务，执行完后还原为最开始的状态，如果当前线程已经设置了traceId就加到前面
     */
    public static void runWithExtraTraceId(String extraTraceId, Runnable runnable) {
        String currentTraceId = MDC.get(ConstantKeys.TAG_TRACE_ID);
        try {
            if (null != extraTraceId && !extraTraceId.isEmpty()) {
                MDC.put(ConstantKeys.TAG_TRACE_ID, extraTraceId + " " + currentTraceId);
            }
            runnable.run();
        } finally {
            MDC.put(ConstantKeys.TAG_TRACE_ID, currentTraceId);
        }
    }

    /**
     * 使用指定的traceId标记任务，执行完后还原为最开始的状态并返回结果，如果当前线程已经设置了traceId就加到前面
     */
    public static <T> T callWithExtraTraceId(String extraTraceId, Callable<T> runnable) throws Exception {
        String currentTraceId = MDC.get(ConstantKeys.TAG_TRACE_ID);
        try {
            if (null != extraTraceId && !extraTraceId.isEmpty()) {
                MDC.put(ConstantKeys.TAG_TRACE_ID, extraTraceId + " " + currentTraceId);
            }
            return runnable.call();
        } finally {
            MDC.put(ConstantKeys.TAG_TRACE_ID, currentTraceId);
        }
    }
}
